"""
名称：CobotGPT_V4_实时停止测试
功能：搭载碰撞检测模型，并实现碰撞检测和实时停止
缺陷：未实现碰撞后轨迹优化
"""
import queue
import time
import torch
from math import pi
from Modules.AUBO.User.Releases.robotcontrol import *
import csv
import timeit
from multiprocessing import Process, Queue

from Modules.KINECT.User.Releases.kinectcontrol import *

import numpy as np
import pandas as pd
import torch
import torch.nn as nn
import matplotlib.pyplot as plt


class DNN(nn.Module):
    def __init__(self):
        """ 搭建神经网络各层 """
        super(DNN, self).__init__()
        self.net = nn.Sequential(  # 按顺序搭建各层
            nn.Linear(36, 32), nn.Sigmoid(),  # 第 1 层：全连接层
            nn.Linear(32, 8), nn.Sigmoid(),  # 第 2 层：全连接层
            nn.Linear(8, 4), nn.Sigmoid(),  # 第 3 层：全连接层
            nn.Linear(4, 1), nn.Sigmoid()  # 第 4 层：全连接层
        )

    def forward(self, x):
        """ 前向传播 """
        y = self.net(x)  # x 即输入数据
        return y  # y 即输出数据


def delay(microseconds):
    """
    功能：延迟xx微秒，1s = 1*10^3 ms = 1*10^6 us = 1*10^9 ns
    输入：
    输出
    """
    start = timeit.default_timer()
    while True:
        elapsed = (timeit.default_timer() - start) * 10 ** 6
        if elapsed >= microseconds:
            break
    return 0



class CobotGPT:

    def __init__(self):
        """
        功能：设定类里面通用的一些属性，静态值，不能放动态的
        输入：
        输出：连接状态码
        """

        # 配置服务器地址和端口
        # self.ip = '169.254.94.55'
        self.ip = '192.168.1.6'  # 一楼：注意配置网线口的适配器属性：关掉IPV6，手动设置IPV4地址和机械臂在同一个网段（README）
        self.port = 8899

        # 预设的机械臂工作路径点：9个
        self.joint_radian_list = [
            ( 80 * pi / 180,   0 * pi / 180,   0 * pi / 180,  0 * pi / 180,  0 * pi / 180, 0 * pi / 180),
            (160 * pi / 180,   0 * pi / 180,   0 * pi / 180,  0 * pi / 180,  0 * pi / 180, 0 * pi / 180),
            (160 * pi / 180, -30 * pi / 180,   0 * pi / 180,  0 * pi / 180,  0 * pi / 180, 0 * pi / 180),
            (160 * pi / 180, -30 * pi / 180,  80 * pi / 180,  0 * pi / 180,  0 * pi / 180, 0 * pi / 180),
            (160 * pi / 180, -30 * pi / 180,  80 * pi / 180, 20 * pi / 180,  0 * pi / 180, 0 * pi / 180),
            (160 * pi / 180, -30 * pi / 180,  80 * pi / 180, 20 * pi / 180, 90 * pi / 180, 0 * pi / 180),  # 第一阶段
            (160 * pi / 180,   0 * pi / 180,  80 * pi / 180, 20 * pi / 180, 90 * pi / 180, 0 * pi / 180),
            ( 90 * pi / 180,   0 * pi / 180,  80 * pi / 180, 20 * pi / 180, 90 * pi / 180, 0 * pi / 180),
            ( 90 * pi / 180,  -20 * pi / 180, 80 * pi / 180, 10 * pi / 180, 90 * pi / 180, 0 * pi / 180)]



    def robot_set_position(self, cnt, queue1: Queue, queue2: Queue):
        """
        功能：设置机械臂的固定轨迹点，并实时获取关节角信息
        输入：cnt输入测试次数，queue1输入是否碰撞，queue2输入机械臂关节角
        输出：
        """

        # 初始化logger
        logger_init()
        # 启动测试
        logger.info("{0} robot_set_position() running...".format(Auboi5Robot.get_local_time()))
        # 系统初始化
        Auboi5Robot.initialize()
        # 创建机械臂控制类
        robot = Auboi5Robot()
        # 创建上下文
        handle = robot.create_context()
        # 打印上下文
        logger.info("Func1：robot.rshd={0}".format(handle))

        try:
            # 连接服务器
            ip = self.ip
            port = self.port
            result = robot.connect(ip, port)
            if result != RobotErrorType.RobotError_SUCC:
                logger.info("connect server{0}:{1} failed.".format(ip, port))
            else:
                # 上电
                robot.robot_startup()
                # 设置碰撞等级
                robot.set_collision_class(10)
                # 初始化全局配置文件
                robot.init_profile()

                # 设置关节最大加速度
                robot.set_joint_maxacc((1.5, 1.5, 1.5, 1.5, 1.5, 1.5))
                # 设置关节最大速度
                # robot.set_joint_maxvelc((0.2, 0.2, 0.2, 0.2, 0.2, 0.2))
                # robot.set_joint_maxvelc((0.1, 0.1, 0.1, 0.1, 0.1, 0.1))
                robot.set_joint_maxvelc((0.3, 0.3, 0.3, 0.3, 0.3, 0.3))

                # 回到初始位姿
                joint_radian_ori = (0, 0, 0, 0, 0, 0)  # 初始位姿,turple类型
                # print(type(joint_radian_ori))
                logger.info("回到初始位姿：move joint to {0}".format(joint_radian_ori))
                robot.move_joint(joint_radian_ori)
                robot.set_joint_maxvelc((0.1, 0.1, 0.1, 0.1, 0.1, 0.1))  # 低速挡
                logger.info("机器人开始工作......")

                # 循环测试
                while result == RobotErrorType.RobotError_SUCC:
                    # while cnt > 0:
                    #     logger.info(f"cnt={cnt}")
                    #     cnt -= 1

                    if queue2.full() is True:
                        try:
                            joint_radian = queue2.get(timeout= 2 )
                            logger.info("move joint to {0}".format(joint_radian))
                            robot.move_joint(joint_radian)
                        except queue.Empty:
                            logger.info("TimeOut: joint_radian_queue or robot_step_queue is empty")
                logger.info("===========robot_set_position fail=========")

        # 出现错误的处理
        except RobotError as e:
            logger.error("{0} robot Event:{1} in robot_set_position()".format(robot.get_local_time(), e))

        except KeyboardInterrupt:
            # 断开服务器链接
            if robot.connected:
                # 断开机械臂链接
                robot.disconnect()
            # 释放库资源
            Auboi5Robot.uninitialize()
            print("robot disconnected-------------------------")

        # 无论是否有异常，都会执行的代码
        finally:
            # 断开服务器链接
            # if robot.connected:
            #     # 关闭机械臂
            #     robot.robot_shutdown()
            #     # 断开机械臂链接
            #     robot.disconnect()
            # # 释放库资源
            # Auboi5Robot.uninitialize()
            logger.info("{0} test completed.".format(Auboi5Robot.get_local_time()))

    def dataset(self, queue1: Queue, queue2: Queue):
        """
        功能：将坐标数据转换成数据集
        输入：queue1:输入骨骼坐标，queue2：输出36列一维数据
        输出：
        """

        # 初始化logger
        logger_init()
        # 启动测试
        logger.info("Func2：{0} dataset() running...".format(Auboi5Robot.get_local_time()))
        # 系统初始化
        Auboi5Robot.initialize()
        # 创建机械臂控制类
        robot = Auboi5Robot()
        # 创建上下文
        handle = robot.create_context()
        # 打印上下文
        logger.info("robot.rshd={0}".format(handle))

        try:
            # 连接服务器
            ip = self.ip
            port = self.port
            result = robot.connect(ip, port)
            if result != RobotErrorType.RobotError_SUCC:
                logger.info("connect server{0}:{1} failed.".format(ip, port))
            else:
                while result == RobotErrorType.RobotError_SUCC:
                    # 获取当前机械臂六个关节角
                    current_waypoint = robot.get_current_waypoint()
                    data1 = current_waypoint["joint"]

                    # 将此刻queue1管道（骨骼坐标）清空，并获取最新骨骼坐标
                    while queue1.empty() is not True:
                        queue1.get()
                    # delay(1000) # 间隔1ms，防止管道来不及填充数据（说明：不能放延迟，会导致程序一卡一卡的）
                    try:
                        data2 = queue1.get(timeout=1)
                    except queue.Empty:
                        data2 = []
                        logger.info("TimeOut: queue is already empty")

                    # queue2管道（36列一维数据）放入数据（36列：骨骼坐标 + 六个关节角）
                    data = data1 + data2

                    # print(f"data:{data}")
                    if queue2.empty() is True:
                        try:
                            queue2.put(data, timeout = 2)
                        except queue.Full:
                            logger.info("TimeOut: queue is already fully")
                    # queue2.put(data)
                logger.info("========dataset fail==========")

        # 出现错误的处理
        except RobotError as e:
            logger.error("{0} robot Event:{1}".format(robot.get_local_time(), e))

    def collision_detection(self, queue1: Queue, queue2: Queue, queue3: Queue):
        """
        功能：碰撞检测
        输入：queue1:输出是否碰撞，queue2:输出下一个关节角，queue3：输入36列一维数据
        输出：
        """

        collision_flag = 0  # 默认未发生碰撞
        step = 0
        model = torch.load('model_collision_detection_v1.pth')  # 加载训练好的模型
        model.to('cuda:0')  # 模型加载到GPU上

        # 初始化logger
        logger_init()
        # 启动测试
        logger.info("{0} collision_detection() running...".format(Auboi5Robot.get_local_time()))
        # 系统初始化
        Auboi5Robot.initialize()
        # 创建机械臂控制类
        robot = Auboi5Robot()
        # 创建上下文
        handle = robot.create_context()
        # 打印上下文
        logger.info("Func3：robot.rshd={0}".format(handle))

        try:
            # 连接服务器
            ip = self.ip
            port = self.port
            result = robot.connect(ip, port)
            if result != RobotErrorType.RobotError_SUCC:
                logger.info("connect server{0}:{1} failed.".format(ip, port))
            else:
                # while True:
                while result == RobotErrorType.RobotError_SUCC:
                    # queue1（是否碰撞）、queue2（下一个关节角）管道战术清空
                    # while queue1.empty() is not True:
                    #     queue1.get()
                    while queue2.empty() is not True:
                        queue2.get()

                    # 运行正常，传输下一个关节角给机械臂
                    if collision_flag == 0:
                        if queue2.empty() is True:
                            try:
                                queue2.put(self.joint_radian_list[step])
                                step += 1
                                if step == int(len(self.joint_radian_list)):
                                    step = 0
                            except queue.Full:
                                logger.info("TimeOut: queue is FULLY")




                        # print("collision 运行正常传输关节角")
                        # queue1.put(0)
                        # queue2.put(self.joint_radian_list[step])
                        # step += 1
                        # if step == int(len(self.joint_radian_list)):
                        #     step = 0
                    # 发生碰撞，传输碰撞信号
                    # else:
                    #     # print("collision 传输碰撞信号")
                    #     # collision_flag = 0
                    #     queue1.put(1)

                    # while queue3.empty() is not True:
                    #     queue3.get()
                    # 碰撞检测模型输入数据获取, 36列
                    try:
                        X = queue3.get(timeout=1)  # 1s的等待时间
                    except queue.Empty:
                        X = []
                        logger.info("TimeOut: position_data_queue is empty")
                    # X = queue3.get()  # 获取36列一维数据



                    # print(f"self.data:{self.data}")
                    # print(len(X))
                    # print(f"X:{X}")

                    if len(X) == 36:
                        arr = np.array(X)  # 列表 => numpy数组
                        arr = arr.astype(np.float32)  # 转为 float32 类型numpy数组
                        ts = torch.tensor(arr)  # numpy数组 => 张量
                        ts = ts.to('cuda')
                        with torch.no_grad():  # 该局部关闭梯度计算功能
                            pred = model(ts)  # 一次前向传播（批量）
                            pred[pred >= 0.5] = 1
                            pred[pred < 0.5] = 0
                            pred = int(pred.item())  # 张量 => int值
                            print(f"是否碰撞: {pred}")
                    else:
                        # print("输入的总坐标格式错误")
                        pred = 0

                    if pred == 1:
                        collision_flag = 1
                        robot.move_pause()
                        # time.sleep(1)
                    else:
                        collision_flag = 0
                        robot.move_continue()
                logger.info("==========collision_detection fail========")

        # 出现错误的处理
        except RobotError as e:
            logger.error("{0} robot Event:{1}".format(robot.get_local_time(), e))



if __name__ == '__main__':
    # 创建类实例
    game = BodyGameRuntime()
    cobot = CobotGPT()

    # 创建通信队列
    q1 = Queue(maxsize=1)  # q1:是否碰撞，
    q2 = Queue(maxsize=1)  # q2:机械臂关节角
    q3 = Queue(maxsize=1)  # q3:kinect_waypoint
    q4 = Queue(maxsize=1)  # q4:36列坐标

    # 创建子进程
    p1 = Process(target=cobot.robot_set_position, kwargs={"cnt": 3, "queue1": q1, "queue2": q2})
    p2 = Process(target=cobot.collision_detection, kwargs={"queue1": q1, "queue2": q2, "queue3": q4})
    p3 = Process(target=cobot.dataset, kwargs={"queue1": q3, "queue2": q4})
    p1.start()
    p2.start()
    p3.start()

    # 主进程
    game.run(q3)
